import pandas
import torch

import kit
from Covid19Dataset import Covid19Dataset
from torch.utils.data import DataLoader
from SimpleModel import SimpleModel

config = {
    'device': 'cpu',
    # 随机数种子
    'seed': 5201314,
    # learning rate
    'lr': 0.8e-4,
    'momentum': 0.9,
    'batch_size': 265,
    'early_stop': 400,
    'epochs': 7000,
    # 将训练集的一些数据分配到验证集
    'valid_ratio': 0.2,
    'save_model_path': './model.ckpt',
    'save_pred_path': './pred.csv',
    'csv_train': './covid.train.csv',
    'csv_test': './covid.test.csv',
}

if __name__ == '__main__':
    # 复现训练
    kit.same_seed(config['seed'])

    # 从csv文件中读取出数据
    csv_train_data = pandas.read_csv(config['csv_train']).values
    csv_test_data = pandas.read_csv(config['csv_test']).values

    # 划分验证集训练集
    csv_train_data, csv_valid_data = kit.train_valid_split(csv_train_data, config['valid_ratio'], seed=config['seed'])

    # 挑选特征和标签
    train_feats, train_labels = kit.train_feats_labels_split(csv_train_data, select_all=False)
    valid_feats, valid_labels = kit.train_feats_labels_split(csv_valid_data, select_all=False)
    test_feats, test_labels = kit.train_feats_labels_split(csv_test_data, True, select_all=False)

    # 建立数据集
    train_dataset = Covid19Dataset(train_feats, train_labels)
    valid_dataset = Covid19Dataset(valid_feats, valid_labels)
    test_dataset = Covid19Dataset(test_feats)

    # 加载数据集
    train_loader = DataLoader(train_dataset, batch_size=config['batch_size'], shuffle=True, pin_memory=True)
    valid_loader = DataLoader(valid_dataset, batch_size=config['batch_size'], shuffle=True, pin_memory=True)
    test_loader = DataLoader(test_dataset, batch_size=config['batch_size'], shuffle=False, pin_memory=True)

    # 创建模型
    model = SimpleModel(input_dim=train_feats.shape[1]).to(config['device'])
    # 训练
    min_loss, model_loss_record = kit.trainer(train_loader, valid_loader, model, config)
    kit.plot_learning_curve(model_loss_record, title='deep model')

    # 预测
    del model
    model = SimpleModel(input_dim=test_feats.shape[1]).to(config['device'])
    ckpt = torch.load(config['save_model_path'], map_location=config['device'])
    model.load_state_dict(ckpt)
    kit.plot_pred(valid_loader, model, config['device'])
    pred_data = kit.tester(test_loader, model, config['device'])  # predict COVID-19 cases with your model
    kit.save_pred(pred_data, config['save_pred_path'])
